{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# For tips on running notebooks in Google Colab, see\n",
        "# https://pytorch.org/tutorials/beginner/colab\n",
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n",
        "# Advanced: Making Dynamic Decisions and the Bi-LSTM CRF\n",
        "\n",
        "## Dynamic versus Static Deep Learning Toolkits\n",
        "\n",
        "Pytorch is a *dynamic* neural network kit. Another example of a dynamic\n",
        "kit is [Dynet](https://github.com/clab/dynet)_ (I mention this because\n",
        "working with Pytorch and Dynet is similar. If you see an example in\n",
        "Dynet, it will probably help you implement it in Pytorch). The opposite\n",
        "is the *static* tool kit, which includes Theano, Keras, TensorFlow, etc.\n",
        "The core difference is the following:\n",
        "\n",
        "* In a static toolkit, you define\n",
        "  a computation graph once, compile it, and then stream instances to it.\n",
        "* In a dynamic toolkit, you define a computation graph *for each\n",
        "  instance*. It is never compiled and is executed on-the-fly\n",
        "\n",
        "Without a lot of experience, it is difficult to appreciate the\n",
        "difference. One example is to suppose we want to build a deep\n",
        "constituent parser. Suppose our model involves roughly the following\n",
        "steps:\n",
        "\n",
        "* We build the tree bottom up\n",
        "* Tag the root nodes (the words of the sentence)\n",
        "* From there, use a neural network and the embeddings\n",
        "  of the words to find combinations that form constituents. Whenever you\n",
        "  form a new constituent, use some sort of technique to get an embedding\n",
        "  of the constituent. In this case, our network architecture will depend\n",
        "  completely on the input sentence. In the sentence \"The green cat\n",
        "  scratched the wall\", at some point in the model, we will want to combine\n",
        "  the span $(i,j,r) = (1, 3, \\text{NP})$ (that is, an NP constituent\n",
        "  spans word 1 to word 3, in this case \"The green cat\").\n",
        "\n",
        "However, another sentence might be \"Somewhere, the big fat cat scratched\n",
        "the wall\". In this sentence, we will want to form the constituent\n",
        "$(2, 4, NP)$ at some point. The constituents we will want to form\n",
        "will depend on the instance. If we just compile the computation graph\n",
        "once, as in a static toolkit, it will be exceptionally difficult or\n",
        "impossible to program this logic. In a dynamic toolkit though, there\n",
        "isn't just 1 pre-defined computation graph. There can be a new\n",
        "computation graph for each instance, so this problem goes away.\n",
        "\n",
        "Dynamic toolkits also have the advantage of being easier to debug and\n",
        "the code more closely resembling the host language (by that I mean that\n",
        "Pytorch and Dynet look more like actual Python code than Keras or\n",
        "Theano).\n",
        "\n",
        "## Bi-LSTM Conditional Random Field Discussion\n",
        "\n",
        "For this section, we will see a full, complicated example of a Bi-LSTM\n",
        "Conditional Random Field for named-entity recognition. The LSTM tagger\n",
        "above is typically sufficient for part-of-speech tagging, but a sequence\n",
        "model like the CRF is really essential for strong performance on NER.\n",
        "Familiarity with CRF's is assumed. Although this name sounds scary, all\n",
        "the model is a CRF but where an LSTM provides the features. This is\n",
        "an advanced model though, far more complicated than any earlier model in\n",
        "this tutorial. If you want to skip it, that is fine. To see if you're\n",
        "ready, see if you can:\n",
        "\n",
        "-  Write the recurrence for the viterbi variable at step i for tag k.\n",
        "-  Modify the above recurrence to compute the forward variables instead.\n",
        "-  Modify again the above recurrence to compute the forward variables in\n",
        "   log-space (hint: log-sum-exp)\n",
        "\n",
        "If you can do those three things, you should be able to understand the\n",
        "code below. Recall that the CRF computes a conditional probability. Let\n",
        "$y$ be a tag sequence and $x$ an input sequence of words.\n",
        "Then we compute\n",
        "\n",
        "$$\\begin{align}P(y|x) = \\frac{\\exp{(\\text{Score}(x, y)})}{\\sum_{y'} \\exp{(\\text{Score}(x, y')})}\\end{align}$$\n",
        "\n",
        "Where the score is determined by defining some log potentials\n",
        "$\\log \\psi_i(x,y)$ such that\n",
        "\n",
        "$$\\begin{align}\\text{Score}(x,y) = \\sum_i \\log \\psi_i(x,y)\\end{align}$$\n",
        "\n",
        "To make the partition function tractable, the potentials must look only\n",
        "at local features.\n",
        "\n",
        "In the Bi-LSTM CRF, we define two kinds of potentials: emission and\n",
        "transition. The emission potential for the word at index $i$ comes\n",
        "from the hidden state of the Bi-LSTM at timestep $i$. The\n",
        "transition scores are stored in a $|T|x|T|$ matrix\n",
        "$\\textbf{P}$, where $T$ is the tag set. In my\n",
        "implementation, $\\textbf{P}_{j,k}$ is the score of transitioning\n",
        "to tag $j$ from tag $k$. So:\n",
        "\n",
        "$$\\begin{align}\\text{Score}(x,y) = \\sum_i \\log \\psi_\\text{EMIT}(y_i \\rightarrow x_i) + \\log \\psi_\\text{TRANS}(y_{i-1} \\rightarrow y_i)\\end{align}$$\n",
        "\n",
        "$$\\begin{align}= \\sum_i h_i[y_i] + \\textbf{P}_{y_i, y_{i-1}}\\end{align}$$\n",
        "\n",
        "where in this second expression, we think of the tags as being assigned\n",
        "unique non-negative indices.\n",
        "\n",
        "If the above discussion was too brief, you can check out\n",
        "[this](http://www.cs.columbia.edu/%7Emcollins/crf.pdf)_ write up from\n",
        "Michael Collins on CRFs.\n",
        "\n",
        "## Implementation Notes\n",
        "\n",
        "The example below implements the forward algorithm in log space to\n",
        "compute the partition function, and the viterbi algorithm to decode.\n",
        "Backpropagation will compute the gradients automatically for us. We\n",
        "don't have to do anything by hand.\n",
        "\n",
        "The implementation is not optimized. If you understand what is going on,\n",
        "you'll probably quickly see that iterating over the next tag in the\n",
        "forward algorithm could probably be done in one big operation. I wanted\n",
        "to code to be more readable. If you want to make the relevant change,\n",
        "you could probably use this tagger for real tasks.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "c:\\Users\\Administrator\\AppData\\Local\\Programs\\Python\\Python38\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
            "  from .autonotebook import tqdm as notebook_tqdm\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "<torch._C.Generator at 0x6f945b0>"
            ]
          },
          "execution_count": 1,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Author: Robert Guthrie\n",
        "\n",
        "import torch\n",
        "import torch.autograd as autograd\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "\n",
        "torch.manual_seed(1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Helper functions to make the code more readable.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def argmax(vec):\n",
        "    # return the argmax as a python int\n",
        "    _, idx = torch.max(vec, 1)\n",
        "    return idx.item()\n",
        "\n",
        "\n",
        "def prepare_sequence(seq, to_ix):\n",
        "    idxs = [to_ix[w] for w in seq]\n",
        "    return torch.tensor(idxs, dtype=torch.long)\n",
        "\n",
        "\n",
        "# Compute log sum exp in a numerically stable way for the forward algorithm\n",
        "def log_sum_exp(vec):\n",
        "    max_score = vec[0, argmax(vec)]\n",
        "    max_score_broadcast = max_score.view(1, -1).expand(1, vec.size()[1])\n",
        "    return max_score + \\\n",
        "        torch.log(torch.sum(torch.exp(vec - max_score_broadcast)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Create model\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "class BiLSTM_CRF(nn.Module):\n",
        "\n",
        "    def __init__(self, vocab_size, tag_to_ix, embedding_dim, hidden_dim):\n",
        "        super(BiLSTM_CRF, self).__init__()\n",
        "        self.embedding_dim = embedding_dim\n",
        "        self.hidden_dim = hidden_dim\n",
        "        self.vocab_size = vocab_size\n",
        "        self.tag_to_ix = tag_to_ix\n",
        "        self.tagset_size = len(tag_to_ix)\n",
        "\n",
        "        self.word_embeds = nn.Embedding(vocab_size, embedding_dim)\n",
        "        self.lstm = nn.LSTM(embedding_dim, hidden_dim // 2,\n",
        "                            num_layers=1, bidirectional=True)\n",
        "\n",
        "        # Maps the output of the LSTM into tag space.\n",
        "        self.hidden2tag = nn.Linear(hidden_dim, self.tagset_size)\n",
        "\n",
        "        # Matrix of transition parameters.  Entry i,j is the score of\n",
        "        # transitioning *to* i *from* j.\n",
        "        self.transitions = nn.Parameter(\n",
        "            torch.randn(self.tagset_size, self.tagset_size))\n",
        "\n",
        "        # These two statements enforce the constraint that we never transfer\n",
        "        # to the start tag and we never transfer from the stop tag\n",
        "        self.transitions.data[tag_to_ix[START_TAG], :] = -10000\n",
        "        self.transitions.data[:, tag_to_ix[STOP_TAG]] = -10000\n",
        "\n",
        "        self.hidden = self.init_hidden()\n",
        "\n",
        "    def init_hidden(self):\n",
        "        return (torch.randn(2, 1, self.hidden_dim // 2),\n",
        "                torch.randn(2, 1, self.hidden_dim // 2))\n",
        "\n",
        "    def _forward_alg(self, feats):\n",
        "        # Do the forward algorithm to compute the partition function\n",
        "        init_alphas = torch.full((1, self.tagset_size), -10000.)\n",
        "        # START_TAG has all of the score.\n",
        "        init_alphas[0][self.tag_to_ix[START_TAG]] = 0.\n",
        "\n",
        "        # Wrap in a variable so that we will get automatic backprop\n",
        "        forward_var = init_alphas\n",
        "\n",
        "        # Iterate through the sentence\n",
        "        for feat in feats:\n",
        "            alphas_t = []  # The forward tensors at this timestep\n",
        "            for next_tag in range(self.tagset_size):\n",
        "                # broadcast the emission score: it is the same regardless of\n",
        "                # the previous tag\n",
        "                emit_score = feat[next_tag].view(\n",
        "                    1, -1).expand(1, self.tagset_size)\n",
        "                # the ith entry of trans_score is the score of transitioning to\n",
        "                # next_tag from i\n",
        "                trans_score = self.transitions[next_tag].view(1, -1)\n",
        "                # The ith entry of next_tag_var is the value for the\n",
        "                # edge (i -> next_tag) before we do log-sum-exp\n",
        "                next_tag_var = forward_var + trans_score + emit_score\n",
        "                # The forward variable for this tag is log-sum-exp of all the\n",
        "                # scores.\n",
        "                alphas_t.append(log_sum_exp(next_tag_var).view(1))\n",
        "            forward_var = torch.cat(alphas_t).view(1, -1)\n",
        "        terminal_var = forward_var + self.transitions[self.tag_to_ix[STOP_TAG]]\n",
        "        alpha = log_sum_exp(terminal_var)\n",
        "        return alpha\n",
        "\n",
        "    def _get_lstm_features(self, sentence):\n",
        "        self.hidden = self.init_hidden()\n",
        "        embeds = self.word_embeds(sentence).view(len(sentence), 1, -1)\n",
        "        lstm_out, self.hidden = self.lstm(embeds, self.hidden)\n",
        "        lstm_out = lstm_out.view(len(sentence), self.hidden_dim)\n",
        "        lstm_feats = self.hidden2tag(lstm_out)\n",
        "        return lstm_feats\n",
        "\n",
        "    def _score_sentence(self, feats, tags):\n",
        "        # Gives the score of a provided tag sequence\n",
        "        score = torch.zeros(1)\n",
        "        tags = torch.cat([torch.tensor([self.tag_to_ix[START_TAG]], dtype=torch.long), tags])\n",
        "        for i, feat in enumerate(feats):\n",
        "            score = score + \\\n",
        "                self.transitions[tags[i + 1], tags[i]] + feat[tags[i + 1]]\n",
        "        score = score + self.transitions[self.tag_to_ix[STOP_TAG], tags[-1]]\n",
        "        return score\n",
        "\n",
        "    def _viterbi_decode(self, feats):\n",
        "        backpointers = []\n",
        "\n",
        "        # Initialize the viterbi variables in log space\n",
        "        init_vvars = torch.full((1, self.tagset_size), -10000.)\n",
        "        init_vvars[0][self.tag_to_ix[START_TAG]] = 0\n",
        "\n",
        "        # forward_var at step i holds the viterbi variables for step i-1\n",
        "        forward_var = init_vvars\n",
        "        for feat in feats:\n",
        "            bptrs_t = []  # holds the backpointers for this step\n",
        "            viterbivars_t = []  # holds the viterbi variables for this step\n",
        "\n",
        "            for next_tag in range(self.tagset_size):\n",
        "                # next_tag_var[i] holds the viterbi variable for tag i at the\n",
        "                # previous step, plus the score of transitioning\n",
        "                # from tag i to next_tag.\n",
        "                # We don't include the emission scores here because the max\n",
        "                # does not depend on them (we add them in below)\n",
        "                next_tag_var = forward_var + self.transitions[next_tag]\n",
        "                best_tag_id = argmax(next_tag_var)\n",
        "                bptrs_t.append(best_tag_id)\n",
        "                viterbivars_t.append(next_tag_var[0][best_tag_id].view(1))\n",
        "            # Now add in the emission scores, and assign forward_var to the set\n",
        "            # of viterbi variables we just computed\n",
        "            forward_var = (torch.cat(viterbivars_t) + feat).view(1, -1)\n",
        "            backpointers.append(bptrs_t)\n",
        "\n",
        "        # Transition to STOP_TAG\n",
        "        terminal_var = forward_var + self.transitions[self.tag_to_ix[STOP_TAG]]\n",
        "        best_tag_id = argmax(terminal_var)\n",
        "        path_score = terminal_var[0][best_tag_id]\n",
        "\n",
        "        # Follow the back pointers to decode the best path.\n",
        "        best_path = [best_tag_id]\n",
        "        for bptrs_t in reversed(backpointers):\n",
        "            best_tag_id = bptrs_t[best_tag_id]\n",
        "            best_path.append(best_tag_id)\n",
        "        # Pop off the start tag (we dont want to return that to the caller)\n",
        "        start = best_path.pop()\n",
        "        assert start == self.tag_to_ix[START_TAG]  # Sanity check\n",
        "        best_path.reverse()\n",
        "        return path_score, best_path\n",
        "\n",
        "    def neg_log_likelihood(self, sentence, tags):\n",
        "        feats = self._get_lstm_features(sentence)\n",
        "        forward_score = self._forward_alg(feats)\n",
        "        gold_score = self._score_sentence(feats, tags)\n",
        "        return forward_score - gold_score\n",
        "\n",
        "    def forward(self, sentence):  # dont confuse this with _forward_alg above.\n",
        "        # Get the emission scores from the BiLSTM\n",
        "        lstm_feats = self._get_lstm_features(sentence)\n",
        "\n",
        "        # Find the best path, given the features.\n",
        "        score, tag_seq = self._viterbi_decode(lstm_feats)\n",
        "        return score, tag_seq"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Run training\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "(tensor(2.6907), [1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1])\n",
            "(tensor(20.4906), [0, 1, 1, 1, 2, 2, 2, 0, 1, 2, 2])\n"
          ]
        }
      ],
      "source": [
        "START_TAG = \"<START>\"\n",
        "STOP_TAG = \"<STOP>\"\n",
        "EMBEDDING_DIM = 5\n",
        "HIDDEN_DIM = 4\n",
        "\n",
        "# Make up some training data\n",
        "training_data = [(\n",
        "    \"the wall street journal reported today that apple corporation made money\".split(),\n",
        "    \"B I I I O O O B I O O\".split()\n",
        "), (\n",
        "    \"georgia tech is a university in georgia\".split(),\n",
        "    \"B I O O O O B\".split()\n",
        ")]\n",
        "\n",
        "word_to_ix = {}\n",
        "for sentence, tags in training_data:\n",
        "    for word in sentence:\n",
        "        if word not in word_to_ix:\n",
        "            word_to_ix[word] = len(word_to_ix)\n",
        "\n",
        "tag_to_ix = {\"B\": 0, \"I\": 1, \"O\": 2, START_TAG: 3, STOP_TAG: 4}\n",
        "\n",
        "model = BiLSTM_CRF(len(word_to_ix), tag_to_ix, EMBEDDING_DIM, HIDDEN_DIM)\n",
        "optimizer = optim.SGD(model.parameters(), lr=0.01, weight_decay=1e-4)\n",
        "\n",
        "# Check predictions before training\n",
        "with torch.no_grad():\n",
        "    precheck_sent = prepare_sequence(training_data[0][0], word_to_ix)\n",
        "    precheck_tags = torch.tensor([tag_to_ix[t] for t in training_data[0][1]], dtype=torch.long)\n",
        "    print(model(precheck_sent))\n",
        "\n",
        "# Make sure prepare_sequence from earlier in the LSTM section is loaded\n",
        "for epoch in range(\n",
        "        300):  # again, normally you would NOT do 300 epochs, it is toy data\n",
        "    for sentence, tags in training_data:\n",
        "        # Step 1. Remember that Pytorch accumulates gradients.\n",
        "        # We need to clear them out before each instance\n",
        "        model.zero_grad()\n",
        "\n",
        "        # Step 2. Get our inputs ready for the network, that is,\n",
        "        # turn them into Tensors of word indices.\n",
        "        sentence_in = prepare_sequence(sentence, word_to_ix)\n",
        "        targets = torch.tensor([tag_to_ix[t] for t in tags], dtype=torch.long)\n",
        "\n",
        "        # Step 3. Run our forward pass.\n",
        "        loss = model.neg_log_likelihood(sentence_in, targets)\n",
        "\n",
        "        # Step 4. Compute the loss, gradients, and update the parameters by\n",
        "        # calling optimizer.step()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "# Check predictions after training\n",
        "with torch.no_grad():\n",
        "    precheck_sent = prepare_sequence(training_data[0][0], word_to_ix)\n",
        "    print(model(precheck_sent))\n",
        "# We got it!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Exercise: A new loss function for discriminative tagging\n",
        "\n",
        "It wasn't really necessary for us to create a computation graph when\n",
        "doing decoding, since we do not backpropagate from the viterbi path\n",
        "score. Since we have it anyway, try training the tagger where the loss\n",
        "function is the difference between the Viterbi path score and the score\n",
        "of the gold-standard path. It should be clear that this function is\n",
        "non-negative and 0 when the predicted tag sequence is the correct tag\n",
        "sequence. This is essentially *structured perceptron*.\n",
        "\n",
        "This modification should be short, since Viterbi and score\\_sentence are\n",
        "already implemented. This is an example of the shape of the computation\n",
        "graph *depending on the training instance*. Although I haven't tried\n",
        "implementing this in a static toolkit, I imagine that it is possible but\n",
        "much less straightforward.\n",
        "\n",
        "Pick up some real data and do a comparison!\n",
        "\n",
        "\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3.8.10 64-bit",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.10"
    },
    "vscode": {
      "interpreter": {
        "hash": "e0144baad0ecee903f108a3e46e51ceadd7da3fc904cfa79747d813b61464b4e"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
